// Copyright (C) Kumo inc. and its affiliates.
// Author: Jeff.li lijippy@163.com
// All rights reserved.
// This program is free software: you can redistribute it and/or modify
// it under the terms of the GNU Affero General Public License as published
// by the Free Software Foundation, either version 3 of the License, or
// (at your option) any later version.
//
// This program is distributed in the hope that it will be useful,
// but WITHOUT ANY WARRANTY; without even the implied warranty of
// MERCHANTABILITY or FITNESS FOR A PARTICULAR PURPOSE.  See the
// GNU Affero General Public License for more details.
//
// You should have received a copy of the GNU Affero General Public License
// along with this program.  If not, see <https://www.gnu.org/licenses/>.
//
#pragma once

#include <gtest/gtest.h>

#include <pollux/testing/gtest_utils.h>
#include <pollux/core/expressions.h>
#include <pollux/core/i_typed_expr.h>
#include <pollux/vector/vector_builder.h>
#include <pollux/testing/functions/prestosql/function_base_test.h>
#include <tests/vector/testing_dictionary_function.h>

namespace kumo::pollux::functions::test {
    using namespace kumo::pollux::test;

    class CastBaseTest : public FunctionBaseTest {
    protected:
        CastBaseTest() {
            exec::registerVectorFunction(
                "testing_dictionary",
                test::TestingDictionaryFunction::signatures(),
                std::make_unique<test::TestingDictionaryFunction>(),
                exec::VectorFunctionMetadataBuilder()
                .defaultNullBehavior(false)
                .build());
        }

        // Build an ITypedExpr for cast(fromType as toType).
        core::TypedExprPtr buildCastExpr(
            const TypePtr &fromType,
            const TypePtr &toType,
            bool isTryCast) {
            core::TypedExprPtr inputField =
                    std::make_shared<const core::FieldAccessTypedExpr>(fromType, "c0");
            return std::make_shared<const core::CastTypedExpr>(
                toType, inputField, isTryCast);
        }

        // Evaluate cast(fromType as toType) and return the result vector.
        VectorPtr evaluateCast(
            const TypePtr &fromType,
            const TypePtr &toType,
            const RowVectorPtr &input,
            bool isTryCast = false) {
            auto castExpr = buildCastExpr(fromType, toType, isTryCast);
            exec::ExprSet exprSet({castExpr}, &execCtx_);
            exec::EvalCtx context(&execCtx_, &exprSet, input.get());

            std::vector<VectorPtr> result(1);
            SelectivityVector rows(input->size());
            exprSet.eval(rows, context, result);
            EXPECT_FALSE(context.errors());
            return result[0];
        }

        // Evaluate cast(fromType as toType) and verify the result matches the
        // expected one.
        void evaluateAndVerify(
            const TypePtr &fromType,
            const TypePtr &toType,
            const RowVectorPtr &input,
            const VectorPtr &expected,
            bool isTryCast = false) {
            auto result = evaluateCast(fromType, toType, input, isTryCast);
            assertEqualVectors(expected, result);
        }

        // Build an ITypedExpr for cast(testing_dictionary(fromType) as toType).
        core::TypedExprPtr buildCastExprWithDictionaryInput(
            const TypePtr &fromType,
            const TypePtr &toType,
            bool isTryCast) {
            core::TypedExprPtr inputField =
                    std::make_shared<const core::FieldAccessTypedExpr>(fromType, "c0");

            // It is not sufficient to wrap input in a dictionary as it will be peeled
            // off before calling "cast". Apply testing_dictionary function to input to
            // ensure that "cast" receives dictionary input.
            core::TypedExprPtr callExpr = std::make_shared<const core::CallTypedExpr>(
                fromType,
                std::vector<core::TypedExprPtr>{inputField},
                "testing_dictionary");
            return std::make_shared<const core::CastTypedExpr>(
                toType, callExpr, isTryCast);
        }

        // Evaluate cast(testing_dictionary(fromType) as toType) and verify the result
        // matches the expected one. Values in expected should correspond to values in
        // input at the same rows.
        void evaluateAndVerifyDictEncoding(
            const TypePtr &fromType,
            const TypePtr &toType,
            const RowVectorPtr &input,
            const VectorPtr &expected,
            bool isTryCast = false) {
            auto castExpr =
                    buildCastExprWithDictionaryInput(fromType, toType, isTryCast);

            VectorPtr result;
            result = evaluate(castExpr, input);

            auto indices = ::kumo::pollux::make_indices_in_reverse(expected->size(), pool());
            assertEqualVectors(wrap_in_dictionary(indices, expected), result);
        }

        // Evaluate try(cast(testing_dictionary(fromType) as toType)) and verify the
        // result matches the expected one. Values in expected should correspond to
        // values in input at the same rows.
        void evaluateAndVerifyCastInTryDictEncoding(
            const TypePtr &fromType,
            const TypePtr &toType,
            const RowVectorPtr &input,
            const VectorPtr &expected) {
            auto castExpr = buildCastExprWithDictionaryInput(fromType, toType, false);
            core::TypedExprPtr tryExpr = std::make_shared<const core::CallTypedExpr>(
                toType, std::vector<core::TypedExprPtr>{castExpr}, "try");

            auto result = evaluate(tryExpr, input);
            auto indices = ::kumo::pollux::make_indices_in_reverse(expected->size(), pool());
            assertEqualVectors(wrap_in_dictionary(indices, expected), result);
        }

        /**
         * @tparam From Source type for cast.
         * @tparam To Destination type for cast.
         * @param typeString Cast type in string.
         * @param input Input vector of type From.
         * @param expectedResult Expected output vector of type To.
         */
        template<typename TFrom, typename TTo>
        void testCast(
            const std::string &typeString,
            const std::vector<std::optional<TFrom> > &input,
            const std::vector<std::optional<TTo> > &expectedResult,
            const TypePtr &fromType = CppToType<TFrom>::create(),
            const TypePtr &toType = CppToType<TTo>::create()) {
            auto result = evaluate(
                fmt::format("cast(c0 as {})", typeString),
                make_row_vector({make_nullable_flat_vector(input, fromType)}));
            auto expected = make_nullable_flat_vector<TTo>(expectedResult, toType);
            assertEqualVectors(expected, result);
        }

        /**
         * @tparam From Source type for cast.
         * @tparam To Destination type for cast.
         * @param typeString Cast type in string.
         * @param input Input vector of type From.
         * @param expectedResult Expected output vector of type To.
         */
        template<typename TFrom, typename TTo>
        void testTryCast(
            const std::string &typeString,
            const std::vector<std::optional<TFrom> > &input,
            const std::vector<std::optional<TTo> > &expectedResult,
            const TypePtr &fromType = CppToType<TFrom>::create(),
            const TypePtr &toType = CppToType<TTo>::create()) {
            auto result = evaluate(
                fmt::format("try_cast(c0 as {})", typeString),
                make_row_vector({make_nullable_flat_vector(input, fromType)}));
            auto expected = make_nullable_flat_vector<TTo>(expectedResult, toType);
            assertEqualVectors(expected, result);
        }

        /**
         * @tparam From Source type for cast.
         * @param typeString Cast type in string.
         * @param input Input vector of type From.
         */
        template<typename TFrom>
        void testInvalidCast(
            const std::string &typeString,
            const std::vector<std::optional<TFrom> > &input,
            const std::string &expectedErrorMessage,
            const TypePtr &fromType = CppToType<TFrom>::create(),
            const bool isRunTimeThrow = false) {
            if (!isRunTimeThrow) {
                POLLUX_ASSERT_USER_THROW(
                    evaluate(
                        fmt::format("cast(c0 as {})", typeString),
                        make_row_vector({make_nullable_flat_vector(input, fromType)})),
                    expectedErrorMessage);
            } else {
                POLLUX_ASSERT_RUNTIME_THROW(
                    evaluate(
                        fmt::format("cast(c0 as {})", typeString),
                        make_row_vector({make_nullable_flat_vector(input, fromType)})),
                    expectedErrorMessage);
            }
        }

        void testCast(
            const VectorPtr &input,
            const VectorPtr &expected,
            std::optional<bool> isTryCast = std::nullopt) {
            const auto &fromType = input->type();
            const auto &toType = expected->type();
            SCOPED_TRACE(fmt::format(
                "Cast from {} to {}", fromType->toString(), toType->toString()));
            const auto copy = createCopy(input);
            // Test with flat encoding.
            {
                SCOPED_TRACE("Flat encoding");
                if (isTryCast.has_value()) {
                    evaluateAndVerify(
                        fromType,
                        toType,
                        make_row_vector({input}),
                        expected,
                        isTryCast.value());
                } else {
                    evaluateAndVerify(fromType, toType, make_row_vector({input}), expected);
                    evaluateAndVerify(
                        fromType, toType, make_row_vector({input}), expected, true);
                }

                // Make sure the input vector does not change.
                assertEqualVectors(input, copy);
            }

            // Test with constant encoding that repeats the first element five times.
            {
                SCOPED_TRACE("Constant encoding");
                const auto constant_row =
                        make_row_vector({BaseVector::wrap_in_constant(5, 0, input)});
                const auto localCopy = createCopy(constant_row);
                const auto constExpected = BaseVector::wrap_in_constant(5, 0, expected);

                if (isTryCast.has_value()) {
                    evaluateAndVerify(
                        fromType, toType, constant_row, constExpected, isTryCast.value());
                } else {
                    evaluateAndVerify(fromType, toType, constant_row, constExpected);
                    evaluateAndVerify(fromType, toType, constant_row, constExpected, true);
                }

                // Make sure the input vector does not change.
                assertEqualVectors(constant_row, localCopy);
                assertEqualVectors(input, copy);
            }

            // Test with dictionary encoding that reverses the indices.
            {
                SCOPED_TRACE("Dictionary encoding");
                if (isTryCast.has_value()) {
                    evaluateAndVerifyDictEncoding(
                        fromType,
                        toType,
                        make_row_vector({input}),
                        expected,
                        isTryCast.value());
                } else {
                    evaluateAndVerifyDictEncoding(
                        fromType, toType, make_row_vector({input}), expected);
                    evaluateAndVerifyDictEncoding(
                        fromType, toType, make_row_vector({input}), expected, true);
                }

                // Make sure the input vector does not change.
                assertEqualVectors(input, copy);
            }
        }

        template<typename TFrom, typename TTo>
        void testCast(
            const TypePtr &fromType,
            const TypePtr &toType,
            std::vector<std::optional<TFrom> > input,
            std::vector<std::optional<TTo> > expected) {
            auto inputVector = make_nullable_flat_vector<TFrom>(input, fromType);
            auto expectedVector = make_nullable_flat_vector<TTo>(expected, toType);

            testCast(inputVector, expectedVector);
        }

        template<typename TFrom>
        void testThrow(
            const TypePtr &fromType,
            const TypePtr &toType,
            const std::vector<std::optional<TFrom> > &input,
            const std::string &expectedErrorMessage) {
            POLLUX_ASSERT_THROW(
                evaluateCast(
                    fromType,
                    toType,
                    make_row_vector({make_nullable_flat_vector<TFrom>(input, fromType)})),
                expectedErrorMessage);
        }

        VectorPtr createCopy(const VectorPtr &input) {
            VectorPtr result;
            SelectivityVector rows(input->size());
            BaseVector::ensure_writable(rows, input->type(), input->pool(), result);
            result->copy(input.get(), rows, nullptr);
            return result;
        }

        std::shared_ptr<core::CastTypedExpr> makeCastExpr(
            const core::TypedExprPtr &input,
            const TypePtr &toType,
            bool nullOnFailure) {
            std::vector<core::TypedExprPtr> inputs = {input};
            return std::make_shared<core::CastTypedExpr>(toType, inputs, nullOnFailure);
        }

        const float kInf = std::numeric_limits<float>::infinity();
        const float kNan = std::numeric_limits<float>::quiet_NaN();
    };
} // namespace kumo::pollux::functions::test
